from data_provider.data_loader import ChannelDataset
from torch.utils.data import DataLoader
import torch
from .spatial_corr_loader import SpatialCorrDataset

data_dict = {
    'CHANNEL': ChannelDataset,
}

# 添加缓存字典用于存储已加载的数据集
_dataset_cache = {}

def data_provider(args, flag, use_cache=True):
    """
    根据参数加载指定的数据集
    
    Args:
        args: 包含数据加载参数的命名空间
        flag: 'train'/'val'/'test'/'all' 分别表示训练集、验证集、测试集和所有数据
        use_cache: 是否使用缓存，避免重复加载数据
        
    Returns:
        data: 数据集对象
        data_loader: DataLoader对象
    """
    # 依据选择的数据类型提供对应的数据集
    if args.data == 'CHANNEL':
        # 使用信道特征数据集
        data_set = ChannelDataset(
            root_path=args.root_path,
            data_path=args.data_path,
            flag=flag,
            size=[args.seq_len, args.label_len, args.pred_len],
            features=args.features,
            n_clusters=args.n_clusters,
            loader=args.loader,
            load_spatial_corr=getattr(args, 'load_spatial_corr', True)  # 默认加载空间相关性数据
        )
    elif args.data == 'SPATIAL_CORR':
        # 使用纯空间相关性数据集
        data_set = SpatialCorrDataset(
            root_path=args.root_path,
            data_path=args.data_path, 
            flag=flag,
            size=[args.seq_len, args.label_len, args.pred_len],
            features=args.features
        )
    else:
        raise ValueError(f"未知的数据类型: {args.data}")
    
    # 打印数据集大小
    print(f"{args.data} 数据集已加载, {flag}集大小: {len(data_set)}")
    
    # 创建DataLoader
    if flag == 'train':
        shuffle_flag = True
    else:
        shuffle_flag = False
        
    # 计算有效批次大小，防止最后一个批次样本不足
    batch_size = min(args.batch_size, len(data_set))
    
    # 优化数据加载器配置
    # 确保使用合理的worker数量（避免过多worker导致内存占用过大）
    num_workers = min(args.num_workers, 8)  # 限制最大worker数为8
    if num_workers < 0:
        num_workers = 0
    
    # 创建数据加载器
    data_loader = DataLoader(
        data_set,
        batch_size=batch_size,
        shuffle=shuffle_flag,
        num_workers=num_workers,  # 使用优化后的worker数量
        drop_last=False,  # 保留所有样本
        pin_memory=True,  # 启用内存固定，加速GPU数据传输
        persistent_workers=True if num_workers > 0 else False,  # 持久化worker进程
        prefetch_factor=2 if num_workers > 0 else 2  # 预取因子，每个worker预取2个batch
    )
    
    print(f"数据加载器配置: batch_size={batch_size}, num_workers={num_workers}, pin_memory=True")
    
    return data_set, data_loader 